package com.bah.tract.security;

import java.io.IOException;
import java.security.KeyStoreException;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.UnrecoverableKeyException;
import java.util.Map;

import javax.security.auth.callback.Callback;
import javax.security.auth.callback.UnsupportedCallbackException;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;

import com.sun.xml.wss.AliasSelector;
import com.sun.xml.wss.impl.callback.PrivateKeyCallback;

/**
 * Adds support for a PrivateKeyCallback and AliasSelector to
 * TrustStoreCallbackImpl. Configuration file needs to contain the following
 * properties. keyStore.location - The location of the keystore file. Should be
 * in JKS format; keyStore.password - The keystore's password; keyStore.alias -
 * Alias of the certificate to use as the private key.
 */
public class KeyStoreCallbackImpl extends TrustStoreCallbackImpl implements
		AliasSelector {

	/**
	 * Logger.
	 */
	private static final Log LOG = LogFactory
			.getLog(KeyStoreCallbackImpl.class);

	/**
	 * Default Constructor.
	 */
	public KeyStoreCallbackImpl() {
		super("c:/projects/tract/keystore.jks", "changeit");
	}

	/**
	 * Handles PrivateKeyCallbacks.
	 * 
	 * @param callback
	 *            a PrivateKeyCallback.
	 * @throws IOException
	 *             if anything goes wrong with the keystore.
	 * @throws UnsupportedCallbackException
	 *             if the given callback is not a PrivateKeyCallback.
	 */
	@Override
	protected final void handle(final Callback callback) throws IOException,
			UnsupportedCallbackException {
		try {
			if (callback instanceof PrivateKeyCallback) {
				final String password = "changeit";
				final PrivateKeyCallback pkCallback = (PrivateKeyCallback) callback;
				final PrivateKey key = (PrivateKey) getKeyStore().getKey(
						pkCallback.getAlias(), password.toCharArray());
				pkCallback.setKey(key);
			} else {
				throw new UnsupportedCallbackException(callback);
			}
		} catch (UnrecoverableKeyException e) {
			throw new IOException(e);
		} catch (KeyStoreException e) {
			throw new IOException(e);
		} catch (NoSuchAlgorithmException e) {
			throw new IOException(e);
		}
	}

	/**
	 * Provides a dynamic way to select the certificate used by a web service to
	 * sign messages.
	 * 
	 * @param map
	 *            ignored
	 * @return The alias specified in the keyStore.alias configuration property.
	 */
	@SuppressWarnings("rawtypes")
	public final String select(final Map map) {
		final String alias = "tract";
		LOG.info("Using certificate with alias - " + alias);
		return alias;
	}
}
